{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2M4YKvaiYihH"
   },
   "outputs": [],
   "source": [
    "import importlib\n",
    "import sys\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "# Install drake.\n",
    "if 'google.colab' in sys.modules and importlib.util.find_spec('pydrake') is None:\n",
    "  version='20200918'\n",
    "  build='nightly'\n",
    "  urlretrieve(f\"https://drake-packages.csail.mit.edu/drake/{build}/drake-{version}/setup_drake_colab.py\", \"setup_drake_colab.py\")\n",
    "  from setup_drake_colab import setup_drake\n",
    "  setup_drake(version=version, build=build)\n",
    "\n",
    "# Determine if this notebook is currently running as a notebook or a unit test.\n",
    "from IPython import get_ipython\n",
    "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.pyplot as plt\n",
    "import open3d\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "from pydrake.all import (RigidTransform, RotationMatrix)\n",
    "\n",
    "\n",
    "def MakeRandomObjectModelAndScenePoints(\n",
    "    num_model_points=20, \n",
    "    noise_std=0, \n",
    "    num_outliers=0, \n",
    "    yaw_O=None,\n",
    "    p_O=None, \n",
    "    num_viewable_points=None, \n",
    "    seed=None):\n",
    "    \"\"\" Returns p_Om, p_s \"\"\"\n",
    "    random_state = np.random.RandomState(seed)\n",
    "\n",
    "    # Make a random set of points to define our object in the x,y plane\n",
    "    theta = np.arange(0, 2.0*np.pi, 2.0*np.pi/num_model_points)\n",
    "    l = 1.0 + 0.5*np.sin(2.0*theta) + 0.4*random_state.rand(1, num_model_points)\n",
    "    p_Om = np.vstack((l * np.sin(theta), l * np.cos(theta), 0 * l))\n",
    "\n",
    "    # Make a random object pose if one is not specified, and apply it to get the scene points.\n",
    "    if p_O is None:\n",
    "        p_O = [2.0*random_state.rand(), 2.0*random_state.rand(), 0.0]\n",
    "    if len(p_O) == 2:\n",
    "        p_O.append(0.0)\n",
    "    if yaw_O is None:\n",
    "        yaw_O = 0.5*random_state.random()\n",
    "    X_O = RigidTransform(RotationMatrix.MakeZRotation(yaw_O), p_O)\n",
    "    if num_viewable_points is None:\n",
    "        num_viewable_points = num_model_points\n",
    "    assert num_viewable_points <= num_model_points\n",
    "    p_s = X_O.multiply(p_Om[:,:num_viewable_points])\n",
    "    p_s[:2, :]  += random_state.normal(scale=noise_std, size=(2, num_viewable_points))\n",
    "    if num_outliers:\n",
    "        outliers = random_state.uniform(low=-1.5, high=3.5, size=(3, num_outliers))\n",
    "        outliers[2,:] = 0\n",
    "        p_s = np.hstack((p_s, outliers))\n",
    "\n",
    "    return p_Om, p_s, X_O\n",
    "\n",
    "def MakeRectangleModelAndScenePoints(\n",
    "    num_points_per_side=7,\n",
    "    noise_std=0, \n",
    "    num_outliers=0, \n",
    "    yaw_O=None,\n",
    "    p_O=None, \n",
    "    num_viewable_points=None, \n",
    "    seed=None):\n",
    "    random_state = np.random.RandomState(seed)\n",
    "    if p_O is None:\n",
    "        p_O = [2.0*random_state.rand(), 2.0*random_state.rand(), 0.0]\n",
    "    if len(p_O) == 2:\n",
    "        p_O.append(0.0)\n",
    "    if yaw_O is None:\n",
    "        yaw_O = 0.5*random_state.random()\n",
    "    X_O = RigidTransform(RotationMatrix.MakeZRotation(yaw_O), p_O)\n",
    "    if num_viewable_points is None:\n",
    "        num_viewable_points = 4*num_points_per_side\n",
    "    \n",
    "    x = np.arange(-1, 1, 2/num_points_per_side)\n",
    "    half_width = 2\n",
    "    half_height = 1\n",
    "    top = np.vstack((half_width*x, half_height + 0*x))\n",
    "    right = np.vstack((half_width + 0*x, -half_height*x))\n",
    "    bottom = np.vstack((-half_width*x, -half_height + 0*x))\n",
    "    left = np.vstack((-half_width + 0*x, half_height*x))\n",
    "    p_Om = np.vstack((np.hstack((top, right, bottom, left)), np.zeros((1, 4*num_points_per_side))))\n",
    "    p_s = X_O.multiply(p_Om[:,:num_viewable_points])\n",
    "    p_s[:2, :]  += random_state.normal(scale=noise_std, size=(2, num_viewable_points))\n",
    "    if num_outliers:\n",
    "        outliers = random_state.uniform(low=-1.5, high=3.5, size=(3, num_outliers))\n",
    "        outliers[2,:] = 0\n",
    "        p_s = np.hstack((p_s, outliers))\n",
    "\n",
    "    return p_Om, p_s, X_O\n",
    "\n",
    "def PlotEstimate(p_Om, p_s, Xhat_O=RigidTransform(), chat=None, X_O=None, ax=None):\n",
    "    p_m = Xhat_O.multiply(p_Om)\n",
    "    if ax is None:\n",
    "        ax = plt.subplot()\n",
    "    Nm = p_Om.shape[1]\n",
    "    artists = ax.plot(p_m[0, :], p_m[1, :], 'bo')\n",
    "    artists += ax.fill(p_m[0, :], p_m[1, :], 'lightblue', alpha=0.5)\n",
    "    artists += ax.plot(p_s[0, :], p_s[1, :], 'ro')\n",
    "    if chat is not None:\n",
    "        artists += ax.plot(np.vstack((p_m[0, chat], p_s[0, :])), np.vstack((p_m[1, chat], p_s[1, :])), 'g--')\n",
    "    if X_O:\n",
    "        p_s = X_O.multiply(p_Om)\n",
    "    artists += ax.fill(p_s[0, :Nm], p_s[1, :Nm], 'lightsalmon')\n",
    "    ax.axis('equal')\n",
    "    return artists\n",
    "\n",
    "def PrintResults(X_O, Xhat_O):\n",
    "    p = X_O.translation()\n",
    "    aa = X_O.rotation().ToAngleAxis()\n",
    "    print(f\"True position: {p}\")\n",
    "    print(f\"True orientation: {aa}\")\n",
    "    p = Xhat_O.translation()\n",
    "    aa = Xhat_O.rotation().ToAngleAxis()\n",
    "    print(f\"Estimated position: {p}\")\n",
    "    print(f\"Estimated orientation: {aa}\")\n",
    "\n",
    "def PoseEstimationGivenCorrespondences(p_Om, p_s, chat):\n",
    "    \"\"\" Returns optimal X_O given the correspondences \"\"\"\n",
    "    # Apply correspondences, and transpose data to support numpy broadcasting\n",
    "    p_Omc = p_Om[:, chat].T\n",
    "    p_s = p_s.T\n",
    "\n",
    "    # Calculate the central points\n",
    "    p_Ombar = p_Omc.mean(axis=0)\n",
    "    p_sbar = p_s.mean(axis=0)\n",
    "\n",
    "    # Calculate the \"error\" terms, and form the data matrix\n",
    "    merr = p_Omc - p_Ombar\n",
    "    serr = p_s - p_sbar\n",
    "    W = np.matmul(serr.T, merr)\n",
    "\n",
    "    # Compute R\n",
    "    U, Sigma, Vt = np.linalg.svd(W)\n",
    "    R = np.matmul(U, Vt)\n",
    "    if np.linalg.det(R) < 0:\n",
    "       print(\"fixing improper rotation\")\n",
    "       Vt[-1, :] *= -1\n",
    "       R = np.matmul(U, Vt)\n",
    "\n",
    "    # Compute p\n",
    "    p = p_sbar - np.matmul(R, p_Ombar)\n",
    "\n",
    "    return RigidTransform(RotationMatrix(R), p)\n",
    "\n",
    "\n",
    "def FindClosestPoints(point_cloud_A, point_cloud_B):\n",
    "    \"\"\"\n",
    "    Finds the nearest (Euclidean) neighbor in point_cloud_B for each\n",
    "    point in point_cloud_A.\n",
    "    @param point_cloud_A A 3xN numpy array of points.\n",
    "    @param point_cloud_B A 3xN numpy array of points.\n",
    "    @return indices An (N, ) numpy array of the indices in point_cloud_B of each\n",
    "        point_cloud_A point's nearest neighbor.\n",
    "    \"\"\"\n",
    "    indices = np.empty(point_cloud_A.shape[1], dtype=int)\n",
    "\n",
    "    kdtree = open3d.geometry.KDTreeFlann(point_cloud_B)\n",
    "    for i in range(point_cloud_A.shape[1]):\n",
    "        nn = kdtree.search_knn_vector_3d(point_cloud_A[:,i], 1)\n",
    "        indices[i] = nn[1][0]\n",
    "\n",
    "    return indices\n",
    "\n",
    "def IterativeClosestPoint(p_Om, p_s, X_O=None, animate=True):\n",
    "    Xhat = RigidTransform()\n",
    "    Nm = p_s.shape[1]\n",
    "    chat_previous = np.zeros(Nm)\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    frames = []\n",
    "    frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=None, X_O=X_O, ax=ax))\n",
    "\n",
    "    while True:\n",
    "        chat = FindClosestPoints(p_s, Xhat.multiply(p_Om))\n",
    "        if np.array_equal(chat, chat_previous):\n",
    "            # Then I've converged.\n",
    "            break\n",
    "        chat_previous = chat\n",
    "        frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=chat, X_O=X_O, ax=ax))\n",
    "        Xhat = PoseEstimationGivenCorrespondences(p_Om, p_s, chat)\n",
    "        frames.append(PlotEstimate(p_Om=p_Om, p_s=p_s, Xhat_O=Xhat, chat=None, X_O=X_O, ax=ax))\n",
    "\n",
    "    ani = animation.ArtistAnimation(fig, frames, interval=400, repeat=False)\n",
    "\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "    plt.close()\n",
    "\n",
    "    if X_O:\n",
    "        PrintResults(X_O, Xhat)\n",
    "\n",
    "    return Xhat, chat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Here is your sandbox.\n",
    "\n",
    "[Here is a set of questions to guide you](https://itempool.com/MIT-Robotic-Manipulation/c/ofuzgx-lLCG)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Svk0qPG_ZRE3",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#p_Om, p_s, X_O = MakeRandomObjectModelAndScenePoints(\n",
    "p_Om, p_s, X_O = MakeRectangleModelAndScenePoints(\n",
    "#    noise_std=0.2, # adds noise to each scene point (default is 0.0)\n",
    "#    num_outliers=3, # adds random points from a uniform distribution   \n",
    "#    num_viewable_points=9, # only this number of model points appear in the scene points\n",
    "    yaw_O = 0.2,  # object orientation (comment it out for random rotations)\n",
    "     p_O = [1, 2],  # object position (comment it out for random positions)\n",
    "    )\n",
    "IterativeClosestPoint(p_Om, p_s, X_O);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Is least-squares the right cost function?\n",
    "\n",
    "Here is a particular setup that is interesting.  The configuration I've given you below results in ICP getting stuck in a local minima.  You will find that the system converges to this local minima from a wide variety of initial conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UZvT8YCQZ2YS",
    "tags": []
   },
   "outputs": [],
   "source": [
    "p_Om, p_s, X_O = MakeRectangleModelAndScenePoints(\n",
    "    num_viewable_points=9, \n",
    "    yaw_O = 0.2,  \n",
    "    p_O = [1, 2], \n",
    ")\n",
    "IterativeClosestPoint(p_Om, p_s, X_O);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Robotic Manipulation -- Lecture 7 Breakout.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
